from ai_lab import AiLabConfig


def cal_model_params(config)->int:
    """
    返回参数以万为单位计算
    """
    # 参数计算
    # 1. 嵌入层（Embedding）和输出层（通常与嵌入层参数共享，若独立计算则为：）
    embedding_params_num = config.vocab_size* config.hidden_size
    in_out_params_num=embedding_params_num*2

    # 2. 注意力层（Multi-Head Attention）
    # 2.1 Q 投影：hidden_size × (num_attention_heads × head_dim)
    q_paras_num= config.hidden_size*config.num_attention_heads*config.head_dim
    # 2.2 K/V 投影：hidden_size × (num_key_value_heads × head_dim × 2)（K 和 V 各一套）
    kv_paras_num= config.hidden_size*config.num_key_value_heads*config.head_dim *2
    # 2.3 输出投影：(num_attention_heads × head_dim) × hidden_size （和q_paras_num 值相同）
    output_paras_num=config.num_attention_heads*config.head_dim*config.hidden_size
    # 2.4 层归一化：每个注意力层配 1 个，参数为 2 × hidden_size
    norm_paras_num= config.hidden_size*2
    atten_paras_num = q_paras_num+kv_paras_num+output_paras_num+norm_paras_num

    # 3. 前馈网络（Feed Forward）
    # 3.1 第一层：hidden_size × intermediate_size1024 × 3072 = 3, 145, 728
    feed_nums1= config.hidden_size*config.intermediate_size
    # 3.2 第二层：intermediate_size × hidden_size3072 × 1024 = 3, 145, 728
    feed_nums2 = config.intermediate_size * config.hidden_size
    # 3.3 偏置项：intermediate_size + hidden_size3072 + 1024 = 4, 0
    feed_abias_nums =config.intermediate_size + config.hidden_size
    # 3.4 层归一化：每个前馈网络配
    feed_norm_paras_num = config.hidden_size * 2
    feed_total_num= feed_nums1+feed_nums2+feed_abias_nums+feed_norm_paras_num

    # 4 隐藏层总参数 (每个隐藏层包含 1 个注意力层 + 1 个前馈网络) * 总层数
    hiden_params_num= config.num_hidden_layers * (atten_paras_num+feed_total_num)

    total_params= in_out_params_num+hiden_params_num
    total_params= int(total_params/(10000))
    print(f'参数数量:{total_params}')
    return total_params